"""
=======================
Visualization utilities
=======================

.. note::
    Try on `Colab <https://colab.research.google.com/github/pytorch/vision/blob/gh-pages/main/_generated_ipynb_notebooks/plot_visualization_utils.ipynb>`_
    or :ref:`go to the end <sphx_glr_download_auto_examples_others_plot_visualization_utils.py>` to download the full example code.

This example illustrates some of the utilities that torchvision offers for
visualizing images, bounding boxes, segmentation masks and keypoints.
"""

# sphinx_gallery_thumbnail_path = "../../gallery/assets/visualization_utils_thumbnail2.png"

import torch
from torch.utils.data import DataLoader
from torchvision import datasets, transforms
from torchvision.utils import make_grid
from torch.autograd import Variable
from MyDataset import MyDataset
import torchvision.transforms.functional as F
import cv2
import numpy as np
import matplotlib.pyplot as plt
import time

plt.rcParams['interactive'] = "true"
plt.rcParams["savefig.bbox"] = 'tight'

def show(imgs):
    if not isinstance(imgs, list):
        imgs = [imgs]
    fig, axs = plt.subplots(ncols=len(imgs), squeeze=False)
    for i, img in enumerate(imgs):
        img = img.detach()
        img = F.to_pil_image(img)
        axs[0, i].imshow(np.asarray(img))
        axs[0, i].set(xticklabels=[], yticklabels=[], xticks=[], yticks=[])

class Model(torch.nn.Module):
    def __init__(self) :
        super(Model, self).__init__()
        self.conv1 = torch.nn.Sequential(torch.nn.Conv2d(1, 64, 3, 1, 1),
                                         torch.nn.ReLU(),
                                         torch.nn.Conv2d(64, 128, 3, 1, 1),
                                         torch.nn.ReLU(),
                                         torch.nn.MaxPool2d(2, 2)) 
        self.dense = torch.nn.Sequential(torch.nn.Linear(14*14*128, 1024),
                                         torch.nn.ReLU(),
                                         torch.nn.Dropout(p = 0.5),
                                         torch.nn.Linear(1024, 11))
        
    def forward(self, x) :
        x = self.conv1(x)
        # x = x.view(-1, 14*14*128)
        x = x.view(-1, 14*14*128)
        x = self.dense(x)
        return x

# 数据预处理
# transform = transforms.Compose([
#     transforms.RandomResizedCrop(28),# 对图像进行随机的crop以后再resize成固定大小
#     transforms.ToTensor(),
#     transforms.Normalize(mean=[0.485, 0.456, 0.406],std=[0.229, 0.224, 0.2]),  # ImageNet全部图片的平均值和标准差
#     transforms.RandomRotation(20), # 随机旋转角度
#     transforms.RandomHorizontalFlip(p=0.5), # 随机水平翻转
# ])
normalize=transforms.Normalize(mean=[.5,.5,.5],std=[.5,.5,.5])
train_transform = transforms.Compose([
  transforms.Grayscale(1),
  transforms.ToTensor(),
#   normalize,
])

train_dataset = datasets.ImageFolder(root= 'data/SM/train4', transform = train_transform)
test_dataset = datasets.MNIST(root = 'data/', train = False, 
                               transform = transforms.ToTensor(), download = False)

# print(train_dataset[0])

# img_data = MyDataset("data/SM/train/")
train_loader = DataLoader(dataset = train_dataset, batch_size=8, shuffle=True)
test_loader = DataLoader(dataset = test_dataset, batch_size= 100, shuffle = True)

## for debug
print(len(train_dataset))
print(len(train_loader))
# oneimg, label = train_dataset[0]
# tmpImg = make_grid(oneimg, nrow = 4)
# tmpImg = tmpImg.numpy().transpose(1, 2, 0)
# cv2.imshow('tmpImg', tmpImg)
# cv2.waitKey(0)
# exit()
images, lables = next(iter(train_loader))
img = make_grid(images, nrow = 2)
img = img.numpy().transpose(1, 2, 0)
print(lables)
cv2.imshow('img', img)
cv2.waitKey(0)

device = torch.device('cpu')
model = Model().to(device)
cost = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.Adam(model.parameters())

# 设置PyTorch使用的线程数(会卡死)
# torch.set_num_threads(14)  # 可以根据需要设置为更高或更低的值

start_time = time.time()
print("start train..., start ts:%f" % (start_time))
epochs = 24
for epoch in range(epochs) :
    # train
    sum_loss = 0.0
    train_correct = 0
    for data in train_loader:
        inputs, lables = data
        inputs, lables = Variable(inputs).cpu(), Variable(lables).cpu()
        optimizer.zero_grad()
        outputs = model(inputs)
        # loss = criterion(outputs, labels)
        loss = cost(outputs, lables)
        loss.backward()
        optimizer.step()
        _, id = torch.max(outputs.data, 1)
        sum_loss += loss.data
        train_correct += torch.sum(id == lables.data)
    print('[%d,%d] loss:%.03f' % (epoch + 1, epochs, sum_loss / len(train_loader)))
    print('        correct:%.03f%%' % (100 * train_correct / len(train_dataset)))

elapsed_time = time.time() - start_time
print(f"ready to test? spend time:{elapsed_time}") 
input()
torch.save(model.state_dict(), "parameter.pkl") #saves
exit()
start_time = time.time()
test_correct = 0
for data in test_loader:
    inputs, lables = data
    inputs, lables = Variable(inputs).cpu(), Variable(lables).cpu()
    outputs = model(inputs)
    _, id = torch.max(outputs.data, 1)
    test_correct += torch.sum(id == lables.data)
print("correct:%.3f%%, spend time:%f" % (100 * test_correct / len(test_dataset), elapsed_time))

# torch.save(model.state_dict(), "parameter.pkl") #save
# model.load_state_dict(torch.load('parameter.pkl')) #load
